package com.atguigu.flink.chapter05.sink;

import org.apache.flink.api.common.eventtime.WatermarkStrategy;
import org.apache.flink.api.common.serialization.SimpleStringEncoder;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.api.connector.source.util.ratelimit.RateLimiterStrategy;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.configuration.MemorySize;
import org.apache.flink.connector.datagen.source.DataGeneratorSource;
import org.apache.flink.connector.datagen.source.GeneratorFunction;
import org.apache.flink.connector.file.sink.FileSink;
import org.apache.flink.core.fs.Path;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.filesystem.OutputFileConfig;
import org.apache.flink.streaming.api.functions.sink.filesystem.bucketassigners.DateTimeBucketAssigner;
import org.apache.flink.streaming.api.functions.sink.filesystem.rollingpolicies.DefaultRollingPolicy;

import java.time.Duration;
import java.time.ZoneId;

/**
 * @Author lzc
 * @Date 2023/6/20 09:15
 */
public class FileSinkDemo {
    public static void main(String[] args) {
        Configuration conf = new Configuration();
        conf.setInteger("rest.port", 2000);
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment(conf);
        env.setParallelism(1);
        
        // 为了让文件完整写入结束, 必须开启 checkpoint
        env.enableCheckpointing(2000);
        
        DataGeneratorSource<String> dataGeneratorSource = new DataGeneratorSource<>(
            new GeneratorFunction<Long, String>() {
                @Override
                public String map(Long value) throws Exception {
                    return "Number:" + value;
                }
            },
            Long.MAX_VALUE,
            RateLimiterStrategy.perSecond(10),
            Types.STRING
        );
        
        
        DataStreamSource<String> stream = env.fromSource(dataGeneratorSource,
                                                         WatermarkStrategy.noWatermarks(),
                                                         "data-source"
        );
        
        FileSink<String> filSink = FileSink
            .<String>forRowFormat(new Path("input/"), new SimpleStringEncoder<>("UTF-8"))
            .withBucketAssigner(new DateTimeBucketAssigner<>("yyyy-MM-dd HH", ZoneId.systemDefault()))
            .withOutputFileConfig(
                OutputFileConfig.builder()
                    .withPartPrefix("atguigu-")
                    .withPartSuffix(".log")
                    .build()
            )
            // 滚动策略
            .withRollingPolicy(
                DefaultRollingPolicy.builder()
                    .withRolloverInterval(Duration.ofSeconds(10))
                    .withMaxPartSize(new MemorySize(1024*1024))
                    .build()
            )
            .build();
        stream.sinkTo(filSink).setParallelism(2);
        
        
        try {
            env.execute();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
}
